#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<assert.h>
#include<stdint.h>
#include<sys/types.h>
#include<fcntl.h>
#include "p2p-message.h"
#include "p2p-file.h"
#include "p2p-engine.h"
#include "p2p-config.h"
#include "p2p-piece.h"
#include "p2p-peer.h"
#include "my_log.h"

/*  An integer array that stores which sockfd has a particular piece. 
    A value of 0 indicates client does not hold the piece and it does not 
    know which peer has the piece. 
    A value of 1 indicates client has the piece.
    Any other integer indicates the sockid corresponding to a 
    peer which has that piece.
*/
uint8_t *piece_map;
uint8_t seed = 0;

uint32_t no_of_pieces; 

int handle_msg(int sockfd, char *recv_buf, int bytes_recvd, peer_list_t *peer_list)
{
    uint8_t id;
    uint32_t len;
    int i, return_val, br;
    block_t block;
    peer_t *peer;
    char *pending_data;

    assert(recv_buf != NULL);
    assert(bytes_recvd > 0);
    assert(sockfd > 2);

    if ((peer = search_peer(peer_list->head, sockfd)) == NULL) {
        printf("[p2p-engine] Unable to find peer\n");
        return -1;
    }

    /*  Handle received messages based on what state the client
        is in with regard to particular peer. In INIT state only "Handshake"
        messages would be handled. In CONNECTED state the following messages
        will be handled: Bitfield, Piece, Request. The "id" 
    */

    // Handle pending data
    if (peer->data_pending) {
        assert(MAXBUFLEN >= bytes_recvd + peer->bytes_recvd);
        pending_data = (char *)malloc(bytes_recvd * sizeof(char));
        memcpy(pending_data, recv_buf, bytes_recvd);
        memcpy(recv_buf, peer->data_recvd, peer->bytes_recvd);
        memcpy(recv_buf + peer->bytes_recvd, pending_data, bytes_recvd);
        bytes_recvd += peer->bytes_recvd;
        
        // Cleanup
        free(peer->data_recvd);
        peer->data_pending = 0;
    }
            
    switch(peer->state) {
        case INIT:
            return_val = handle_handshake_msg(sockfd, recv_buf, bytes_recvd);
            if (return_val == -1) {
                printf("[p2p-engine]: handle_handshake_msg failed\n");
                return -1;
            } else if (return_val == 1) {
                if (send_handshake(sockfd) == -1) {
                    printf("[p2p-engine] initiate_handshake failed for sockfd: %d\n", sockfd);
                }
                
                if (send_bitfield_msg(sockfd) == -1) {
                    printf("[p2p-engine]: send_bitfield_msg failed\n");
                    return -1;
                }
                peer->state = CONNECTED;
            }
            //return HANDSHAKE_MSG_LEN;
            return 4;
        break;
       
        case HS_SENT:
            return_val = handle_handshake_msg(sockfd, recv_buf, bytes_recvd);
            if (return_val == -1) {
                printf("[p2p-engine]: handle_handshake_msg failed\n");
                return -1;
            } else if (return_val == 1) {
                if (send_bitfield_msg(sockfd) == -1) {
                    printf("[p2p-engine]: send_bitfield_msg failed\n");
                    return -1;
                }
                peer->state = CONNECTED;
            }
            return HANDSHAKE_MSG_LEN;
        break;
       
        case CONNECTED:
            len = my_ntohl(recv_buf);
            br = len + LEN_LEN;
            id = recv_buf[LENGTH_LEN];
            
            // Handle different messages
            switch(id) {
                case BITFIELD_MSG:
                    LOG(LOG_INFO, "Bitfield Message received from peer:%d", sockfd);
                    if (handle_bitfield_msg(sockfd, recv_buf, peer, piece_map, no_of_pieces) == -1) {
                        printf("[p2p-engine] handle_bitfield_msg failed\n");
                        return -1;
                    }

                    //  If we are seed return. No need to send Request messages. 
                    if (seed)
                        return br;
                    
                    // If the peer doesn't have any useful piece, return
                    if (!peer->piece_count)
                        return br;

                    peer->data_exchange = 1;
                    
                    if (request_block(sockfd, peer, &block, no_of_pieces) == -1) {
                        printf("[p2p-engine] select_block failed\n");
                        return -1;
                    }
                break;
                
                case HAVE_MSG:
                    //  If we are seed return. No need to send Request messages. 
                    if (seed)
                        return br;
                    
                    LOG(LOG_INFO, "Have Message received from Peer:%d", sockfd);
                    if (handle_have_msg(sockfd, recv_buf + LEN_LEN + ID_LEN, peer, piece_map) == -1) {
                        printf("handle_have_msg failed\n");
                        return -1;
                    }
                  
                    // If the peer doesn't have any useful piece, return
                    if (!peer->piece_count)
                        return br;

                    // Send request only if there is no data exchange with this peer
                    if (!peer->data_exchange) {
                        peer->data_exchange = 1;
                        if (request_block(sockfd, peer, &block, no_of_pieces) == -1) {
                            printf("[p2p-engine] select_block failed\n");
                            return -1;
                        }
                    }
                break;
                
                case REQUEST_MSG:
                    LOG(LOG_INFO, "[p2p-engine] Request Message received from Peer:%d", sockfd);
                    assert(bytes_recvd >= LEN_LEN + ID_LEN + INDEX_LEN + BEGIN_LEN + LENGTH_LEN);
                    
                    if (handle_request_msg(sockfd, recv_buf + LEN_LEN + ID_LEN, &block) == -1) {
                        printf("handle_request_msg failed\n");
                        return -1;
                    }
                    
                    // Check if you have the piece requested
                    if(piece_map[block.index] != 1) {
                        printf("[p2p-engine] Piece requested not found\n");
                        return br;
                    }

                    if (send_piece_msg(sockfd, &block) == -1) {
                        printf("send_piece_msg failed\n");
                        return -1;
                    }
                break;

                case PIECE_MSG:
                    if (bytes_recvd < LEN_LEN + len) {
                        LOG(LOG_INFO, "Entire data not received. Wait for data in next recv");
                        LOG(LOG_INFO, "bytes_recvd:%d bytes_expected:%d", bytes_recvd, LEN_LEN + len);
                        peer->data_recvd = (char *)malloc(bytes_recvd * sizeof(char));
                        memcpy(peer->data_recvd, recv_buf, bytes_recvd);
                        peer->bytes_recvd = bytes_recvd;
                        peer->data_pending = 1;
                        return br;
                    } else {
                        assert(bytes_recvd >= LEN_LEN + len);
                    }
                    
                    LOG(LOG_INFO, "Piece Message received from Peer:%d", sockfd);
                    if (handle_piece_msg(sockfd, recv_buf + LEN_LEN + ID_LEN, &block, len) == -1) {
                        printf("handle_piece_msg failed\n");
                        return -1;
                    }
                    
                    // TO DO: handle if block received is not block requested
                    
                    // Set the corresponding block as downloaded
                    peer->piece->block[block.begin / BLOCK_SIZE] = 1;
                    
                    // Select a block
                    if ((return_val = select_block(peer->piece, &block, no_of_pieces)) == -1) {
                        printf("[p2p-engine] select_block failed\n");
                        return -1;
                    } else if (return_val == 1) {
                        if (verify_piece_integrity(FILE_, HASH_FILE, peer->piece->index, no_of_pieces) == -1) {
                            printf("[p2p-engine] verify_piece_integrity failed. Resetting piece\n");
                            // Unmark the piece
                            reset_piece(peer->piece);
                            piece_map[peer->piece->index] = 0;
                            return -1;
                        }                        
                        
                        printf("[p2p-engine] Piece:%d has been downloaded\n", peer->piece->index);
                        
                        // Broadcast HAVE message corresponding to currently downloaded piece to all interested peers
                        if (broadcast_have_msg(peer_list->head, peer->piece->index) == -1) {
                            LOG(LOG_ERR, "broadcast_have_msg failed");
                            return -1;
                        }

                        del_piece(&(peer->piece_head), peer->piece);
                        peer->piece_count--;
                   
                        // Write piece_map to META file
                        write_file(META_FILE, 0, piece_map, no_of_pieces, O_WRONLY | O_CREAT);
                    
                        if (peer->piece_count == 0) {
                            peer->data_exchange = 0;
                            if (!check_missing_piece()) {
                                LOG(LOG_INFO, "We have downloaded all the pieces!");
                                seed = 1;
                                return br;
                            }
                            printf("[p2p-engine] No more new pieces from peer: %d\n", peer->psfd);
                            return br;
                        }
                        
                        if (request_block(sockfd, peer, &block, no_of_pieces) == -1) {
                            printf("[p2p-engine] select_block failed\n");
                            return -1;
                        }

                        return br;
                    }
                    
                    if (send_request_msg(sockfd, &block) == -1) {
                        printf("send_piece_msg failed\n");
                        return -1;
                    }
                break;

                default:
                    printf("[p2p-engine]: id:default. I shouldn't be here\n");
            }
        break;

        default:
            printf("[p2p-engine]: state:default. I shouldn't be here\n");
    }
   
    return br;
}


int init()
{
    int index, br, i, last_piece_size;
    char *buffer, *file_buffer;
    
    // Get the # of pieces
    if ((br = read_file(HASH_FILE, 0, &buffer, -1)) == -1) {
        printf("[p2p-engine]: Hash file: %s not found\n", HASH_FILE);
        return -1;
    }

    no_of_pieces = br / SHA1_HASH_SIZE;
    
    free(buffer);
    
    // Construct piece map
    piece_map = (uint8_t *)malloc(no_of_pieces * sizeof(uint8_t));

    if ((br = read_file(META_FILE, 0, &buffer, -1)) == -1) {
        printf("[p2p-engine]: Meta file: %s not found\n", META_FILE);
        
        // Initialize piece_map and block_map
        for (i = 0; i < no_of_pieces; i++) {
            piece_map[i] = 0;
        }

        // write piece_map to META_FILE
        write_file(META_FILE, 0, piece_map, no_of_pieces, O_WRONLY | O_CREAT);
        
        // create an empty file as a place holder for the file that is to be downloaded
        file_buffer = (uint8_t *)malloc(PIECE_SIZE * sizeof(uint8_t));
        memset(file_buffer, 0, PIECE_SIZE);
        for (i = 0; i < no_of_pieces - 1; i++)
            write_file(FILE_, i * PIECE_SIZE, file_buffer, PIECE_SIZE, O_RDWR | O_CREAT);
        free(file_buffer);
       
        // Last piece may/may not be multiple of PIECE_SIZE
        if (last_piece_size = FILE_SIZE % PIECE_SIZE) {
            file_buffer = (uint8_t *)malloc(last_piece_size * sizeof(uint8_t));
            memset(file_buffer, 0, last_piece_size);
            write_file(FILE_, i * PIECE_SIZE, file_buffer, last_piece_size, O_RDWR | O_CREAT);
            free(file_buffer);
        }
    
    } else {
        if (br != no_of_pieces) {
            LOG(LOG_ERR, "Meta file info not matching with # of pieces obtained from Hash file\n");
            LOG(LOG_ERR, "Meta file pieces:%d Hash file pieces:%d", br, no_of_pieces);
            return -1;
        }
        memcpy(piece_map, buffer, no_of_pieces);
        free(buffer);
    }
    
    // Check if the client is missing any piece
    return check_missing_piece();
}

int connect_to_peers(const socket_t *local_socket, const socket_t *remote_socket, 
                     fd_set *master, int *fdmax, int remote_socket_count, peer_list_t *peer_list)
{
    int i, protocol_no = 0, backlog = 0, sockfd;
    socket_t new_local_socket;
   
    // Set local port to zero so that a random port is choosen to initiate a new connection.
    new_local_socket = *local_socket;
    new_local_socket.port = 0;

    for (i = 0; i < remote_socket_count; i++) {
        if ((sockfd = create_tcp_socket(AF_INET, protocol_no, &new_local_socket, remote_socket + i, backlog)) == -1) {
            printf("[p2p-engine] create_tcp_socket returned error\n");
            return -1;
        }
        FD_SET(sockfd, master); // add to master set
        if (sockfd > *fdmax) 
            *fdmax = sockfd;

        // Add to peer list
        if (!peer_list->count) {
            peer_list->head = init_peer(sockfd);
            peer_list->tail = peer_list->head;
        } else {
            peer_list->tail = add_peer(peer_list->tail, sockfd);
        }
        peer_list->count++;
    }

    return 0;
}

int send_handshake_to_peers(peer_list_t *peer_list)
{
    int i;
    peer_t *peer;
    
    peer = peer_list->head;
    while (peer != NULL) {
        if (send_handshake(peer->psfd) == -1) {
            printf("[p2p-engine] initiate_handshake failed for sockfd: %d\n", peer->psfd);
        }
        peer->state = HS_SENT;
        peer = peer->next;
    }

    return 0;
}

int check_missing_piece()
{
    int i;

    for (i = 0; i < no_of_pieces; i++) {
        if (!piece_map[i]) {
            printf("[p2p-engine] Some pieces are missing\n");
            return 1;
        }
    }
    seed = 1;
    printf("[p2p-engine] We have all the pieces\n");
    return 0;
}

int cleanup(peer_list_t *peer_list, int sockfd)
{
    peer_t *peer;

    peer = search_peer(peer_list->head, sockfd);
    del_peer(&(peer_list->head), &(peer_list->tail), peer);
    peer_list->count--;
}

int request_block(int sockfd, peer_t *peer, block_t *block, int no_of_pieces)
{
    // Select a random piece to download 
    if ((peer->piece = select_random_piece(piece_map, &(peer->piece_head), &(peer->piece_count))) == NULL) {
        printf("[p2p-engine] No more new pieces from peer: %d\n", peer->psfd);
        return 0;
    }

    // Mark the piece so that the client doesn't attempt to download it from other peers
    piece_map[peer->piece->index] = 1;
    
    // Select a block from the piece
    if (select_block(peer->piece, block, no_of_pieces) == -1) {
        printf("[p2p-engine] select_block failed\n");
        return -1;
    }
    
    if (send_request_msg(sockfd, block) == -1) {
        printf("send_request_msg failed\n");
        return -1;
    }
    return 0;
}

